﻿using System;
using System.Runtime.CompilerServices;
using RayDen.Library.Core;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Engines.MCMC;

namespace RayDen.RayEngine.Samplers
{
    public class PSSMLTSampler : ISampler
    {
        public static int MaxMutations = 100;
        public static float LargeStepProbability = 0.05f;

        public static int PreviewPasses = 1;
        public static int SamplesPerPassX4 = 16;
        public static int MagicValue2 = 1;
        public static int SuperSamplingSize = 4;
        public int Width { get; protected set; }
        public int Height { get; protected set; }

        public int screenWidth, screenHeight, screenStartLine;
        int currentSampleScreenX, currentSampleScreenY, currentSubSampleIndex;
        public int pass;
        public bool lowLatency;
        public int valuesPerSample, initializedSamples = 0;

        private IRandomProvider rnd;

        public int tIdx;
        protected int _currentSampleScreenY;

        protected int currentMutation;
        protected MCMCSamplerState state;
        protected Sample currentSample;

        public PSSMLTSampler()
        {
            rnd =
                new FastRandom();
            //new HaltonSequence();
            TotalSamples = 0L;
            state = MCMCSamplerState.Initialized;
        }

        int ScreenHeight
        {
            get { return screenStartLine + screenHeight; }

        }

        public Sample GetSample(Sample prevSample)
        {
            Sample newSample = null;

            switch (state)
            {
                case MCMCSamplerState.Initialized:
                    
                    if (currentSample == null)
                    {
                        currentSample = GetNewSample(true);
                    }

                    newSample = GetNewSample(false);
                    newSample.sampleData = currentSample.sampleData;
                    initializedSamples++;
                    if (initializedSamples > screenWidth*screenHeight*4)
                    {
                        state = MCMCSamplerState.LargeStep;
                    }

                    return newSample;

                case MCMCSamplerState.LargeStep:
                    currentSample = GetNewSample(true);
                    currentMutation = 0;
                    state = MCMCSamplerState.SmallStep;
                    break;
                case MCMCSamplerState.SmallStep:
                    currentSample = Mutate(prevSample??currentSample, currentMutation++, MaxMutations);
                    if (currentMutation >= MaxMutations)
                    {
                        state = MCMCSamplerState.LargeStep;
                        currentSample = GetNewSample(true);
                        currentMutation = 0;
                    }
                    break;
            }

            if (GetLazyValue(currentSample) < LargeStepProbability)
            {
                state = MCMCSamplerState.LargeStep;
            }

            return currentSample;
        }

        public Sample Mutate(Sample s, int mc, int mcMax = 10000)
        {

            if (mc > mcMax)
            {
                return GetSample(null); ;
            }

            if (mc > 1 && rnd.NextFloat() > 0.5f)
            {
                if (rnd.NextFloat() > 0.5f)
                    s.imageX = Mutate(s.imageX, rnd.NextFloat(), rnd.NextFloat());
                if (rnd.NextFloat() > 0.5f)
                    s.imageY = Mutate(s.imageY, rnd.NextFloat(), rnd.NextFloat());
                //return s;
                //if (rnd.NextFloat() > 0.5f)return;
            }

            for (int i = 0; i < s.sampleData.Length; i++)
            {
                s.sampleData[i] = Mutate(s.sampleData[i], rnd.NextFloat(), rnd.NextFloat());
            }
            TotalSamples++;
            return s;
        }

        public static float Mutate(float value, float u, float v)
        {
            float s1 = 1f / 1024f, s2 = 1f / 64f;
            float dv = (float)(s2 * Math.Exp(-Math.Log(s2 / s1) * u));
            if (v < 0.5)
            {
                value += dv; if (value > 1) value -= 1;
            }
            else
            {
                value -= dv; if (value < 0) value += 1;
            }
            return value;
        }

        public static float Mutate(double value, float u, float v)
        {
            float s1 = 1f / 1024f, s2 = 1f / 64f;
            float dv = (float)(s2 * Math.Exp(-Math.Log(s2 / s1) * u));
            float vt = (float)value;
            if (v < 0.5)
            {
                vt += dv; if (vt > 1) vt -= 1;
            }
            else
            {
                vt -= dv; if (vt < 0) vt += 1;
            }
            return vt;
        }
        private Sample GetNewSample(bool regen)
        {
            TotalSamples++;
            Sample sample = new Sample(this);
            if (regen)
                sample.InitSample(10 * 10);
            int scrX, scrY;

            if (!lowLatency || (pass >= MagicValue2))
            {
                // In order to improve ray coherency
                int stepX = currentSubSampleIndex % SuperSamplingSize;
                int stepY = currentSubSampleIndex / SuperSamplingSize;

                scrX = currentSampleScreenX;
                scrY = currentSampleScreenY;

                currentSubSampleIndex++;
                if (currentSubSampleIndex == (SamplesPerPassX4))
                {
                    currentSubSampleIndex = 0;
                    currentSampleScreenX++;
                    if (currentSampleScreenX >= screenWidth)
                    {
                        currentSampleScreenX = 0;
                        currentSampleScreenY++;

                        if (currentSampleScreenY >= ScreenHeight)
                        {
                            currentSampleScreenY = _currentSampleScreenY;
                            pass += SamplesPerPassX4;
                        }
                    }
                }

                float r1 = (stepX + ((float)rnd.NextDouble())) / (float)SuperSamplingSize - 0.5f;
                float r2 = (stepY + ((float)rnd.NextDouble())) / (float)SuperSamplingSize - 0.5f;
                sample.imageX = scrX + r1;
                sample.imageY = scrY + r2;
                //sample.imageY += screenStartLine;

                //sample.Init(this,pass, valuesPerSample, false);
            }
            else if (pass >= PreviewPasses)
            {
                scrX = currentSampleScreenX;
                scrY = currentSampleScreenY;

                currentSampleScreenX++;
                if (currentSampleScreenX >= screenWidth)
                {
                    currentSampleScreenX = 0;
                    currentSampleScreenY++;

                    if (currentSampleScreenY >= ScreenHeight)
                    {
                        currentSampleScreenY = _currentSampleScreenY;
                        pass++;
                    }
                }

                float r1 = ((float)rnd.NextDouble()) - .5f;
                float r2 = ((float)rnd.NextDouble()) - .5f;
                sample.imageX = scrX + r1;
                sample.imageY = scrY + r2;
                //sample.imageY += screenStartLine;
                //sample.Init(this, pass,valuesPerSample, false);
            }
            else
            {
                // In order to update the screen faster for the first 16 passes
                for (; ; )
                {
                    int stepX = pass % SuperSamplingSize;
                    int stepY = (pass / SuperSamplingSize) % SuperSamplingSize;

                    scrX = currentSampleScreenX * SuperSamplingSize + stepX;
                    scrY = currentSampleScreenY * SuperSamplingSize + stepY;

                    currentSampleScreenX++;
                    if ((currentSampleScreenX * SuperSamplingSize) >= screenWidth)
                    {
                        currentSampleScreenX = 0;
                        currentSampleScreenY++;

                        if ((currentSampleScreenY * SuperSamplingSize) >= ScreenHeight)
                        {
                            currentSampleScreenY = _currentSampleScreenY;
                            pass++;
                        }
                    }

                    // Check if we are inside the screen
                    if ((scrX < screenWidth) && (scrY < ScreenHeight))
                    {
                        // Ok, it is a valid sample
                        break;
                    }
                    else if (pass >= PreviewPasses)
                    {
                        sample = this.GetNewSample(regen);
                        return sample;
                    }
                }
                /*
                float r1 = ((float)rnd.NextDouble()) - .5f;
                float r2 = ((float)rnd.NextDouble()) - .5f;

                sample.imageX = scrX + r1;
                sample.imageY = scrY + r2;*/
                //sample.imageY += screenStartLine;
                //sample.Init(this, pass, valuesPerSample, false);
            }
            return sample;
        }


        #region ISampler Members


        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public float GetLazyValue(Sample sample)
        {
            return rnd.NextFloat();
        }

        public float GetLDValue(Sample sample)
        {
            return rnd.NextFloat();

        }

        public void Init(int width, int height)
        {
            rnd.Reinitialise((int)DateTime.Now.Ticks);
            if (screenWidth == 0 && screenHeight == 0)
            {
                Width = screenWidth = width;
                Height = screenHeight = height;
                //screenWidth = width;
                //screenHeight = height;
            }
            screenStartLine = Math.Max(0, screenStartLine);
            _currentSampleScreenY = Math.Max(0, screenStartLine);

            currentSubSampleIndex = 0;
            pass = 0;
            TotalSamples = 0;
        }

        public int GetPass()
        {
            return pass;
        }

        public long TotalSamples { get; set; }
        public float GetLazyValue(Sample sample, int index)
        {
            return sample.sampleData[unchecked(index)];
        }

        #endregion

        public void Reinitialise(int seed)
        {
            this.rnd.Reinitialise(seed);
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public float NextFloat()
        {
            return rnd.NextFloat();
        }
    }
}